package com.leet.bfs;

import java.util.*;

public class DistanceEqualsK {
    // 存放parent的关系
    private Map<TreeNode,TreeNode> parents = new HashMap<>();
    // bfs的队列
    private Queue<TreeNode> q = new LinkedList<>();
    // visited已经访问过的不再访问
    private Set<TreeNode> visited = new HashSet<>();

    // 获取父节点，用Map存储(自己的父亲TreeNode->TreeNode的father)（写一个递归）
    void getParent(TreeNode root, TreeNode parent){
        if(root == null) return;
        parents.put(root, parent);
        getParent(root.left, root);
        getParent(root.right, root);
    }
    public List<Integer> distanceK(TreeNode root, TreeNode target, int k) {
        // init
        List<Integer> ans = new ArrayList<>();
        parents = new HashMap<>();
        q = new LinkedList<>();
        visited = new HashSet<>();
        q.offer(target);
        int level = 0;
        while(!q.isEmpty()){
            int size = q.size();
            if(level == k){
                while(!q.isEmpty()){
                    ans.add(q.poll().val);
                }
                return ans;
            }
            for(int i = 0; i < size; i++){
                TreeNode front = q.poll();
                if(front.left != null && !visited.contains(front.left)){
                    q.offer(front.left);
                    visited.add(front.left);
                }
                if(front.right != null && !visited.contains(front.right)){
                    q.offer(front.right);
                    visited.add(front.right);
                }
                TreeNode father = parents.get(front);
                if(father!= null && !visited.contains(father)){
                    q.offer(father);
                    visited.add(father);
                }
            }
            level++;
        }
            return ans;
    }
}
